import sys
import os
import ctypes
import subprocess
import tkinter as tk
import tkinter.font as tkfont
from tkinter import messagebox
from typing import Optional, List, Tuple
from tkinter import ttk

try:
    import winreg
except ImportError:
    winreg = None


RDP_REG_PATHS = [
    r"SYSTEM\CurrentControlSet\Control\Terminal Server\WinStations\RDP-Tcp",
    # 兼容旧版本（可能不存在）
    r"SYSTEM\CurrentControlSet\Control\Terminal Server\WdS\rdpwd\Tds\tcp",
]
RDP_REG_VALUE = "PortNumber"


def is_admin() -> bool:
    try:
        return ctypes.windll.shell32.IsUserAnAdmin() != 0
    except Exception:
        return False


def relaunch_as_admin(extra_args: List[str]) -> None:
    """以管理员权限重新运行当前脚本并传递参数。"""
    python_exe = sys.executable
    script_path = os.path.abspath(sys.argv[0])
    params = "\"" + script_path + "\" " + " ".join(extra_args)
    # SW_SHOWNORMAL = 1
    ctypes.windll.shell32.ShellExecuteW(None, "runas", python_exe, params, None, 1)


def _read_rdp_port_via_reg_query() -> Optional[int]:
    """通过 reg.exe 查询，兼容某些系统对 winreg 的读取限制。"""
    for path in RDP_REG_PATHS:
        try:
            proc = subprocess.run(
                ["reg", "query", f"HKLM\\{path}", "/v", RDP_REG_VALUE],
                capture_output=True,
                text=True,
            )
            if proc.returncode == 0 and proc.stdout:
                # 解析输出行，例如：
                #    PortNumber    REG_DWORD    0x00000d3d
                for line in proc.stdout.splitlines():
                    if "PortNumber" in line and "REG_DWORD" in line:
                        parts = [p for p in line.strip().split() if p]
                        # 最后一个部分通常是数值，可为十六进制格式（0x...）或十进制
                        try:
                            val_str = parts[-1]
                            # int(val_str, 0) 支持 0x 前缀
                            return int(val_str, 0)
                        except Exception:
                            continue
        except Exception:
            continue
    return None


def read_rdp_port() -> Optional[int]:
    # 先尝试使用 winreg 读取
    if winreg is not None:
        for path in RDP_REG_PATHS:
            try:
                key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, path, 0, winreg.KEY_READ)
                value, regtype = winreg.QueryValueEx(key, RDP_REG_VALUE)
                winreg.CloseKey(key)
                if isinstance(value, int):
                    return value
            except FileNotFoundError:
                continue
            except PermissionError:
                # 权限受限则使用 reg 查询作为回退
                break
            except OSError:
                continue
    # 回退：使用 reg.exe 查询
    return _read_rdp_port_via_reg_query()


def write_rdp_port(port: int) -> None:
    if winreg is None:
        raise RuntimeError("当前环境不支持 winreg，无法写入注册表")
    last_error = None
    for path in RDP_REG_PATHS:
        try:
            key = winreg.OpenKey(
                winreg.HKEY_LOCAL_MACHINE,
                path,
                0,
                winreg.KEY_SET_VALUE | winreg.KEY_WRITE,
            )
            # PortNumber 为 DWORD（REG_DWORD）
            winreg.SetValueEx(key, RDP_REG_VALUE, 0, winreg.REG_DWORD, int(port))
            winreg.CloseKey(key)
        except FileNotFoundError as e:
            last_error = e
            # 某些路径可能不存在，跳过
            continue
        except PermissionError as e:
            last_error = e
            raise
        except OSError as e:
            last_error = e
            raise
    # 如果所有路径都不存在，也视为写入失败
    if last_error and isinstance(last_error, FileNotFoundError):
        # 部分系统只有第一个路径有效，如果它也不存在，提示用户系统不支持或 RDP 未安装
        raise RuntimeError("未找到 RDP 端口配置路径，可能系统未启用远程桌面")


# 已移除防火墙规则自动配置，专注于端口写入与服务重启


def restart_rdp_service() -> Tuple[subprocess.CompletedProcess, Optional[subprocess.CompletedProcess]]:
    # 警告：这将断开当前 RDP 会话
    stop_proc = subprocess.run(["sc", "stop", "TermService"], capture_output=True, text=True)
    start_proc: Optional[subprocess.CompletedProcess] = None

    if stop_proc.returncode == 0:
        # 正常停止后再启动
        try:
            start_proc = subprocess.run(["sc", "start", "TermService"], capture_output=True, text=True)
        except Exception:
            start_proc = None
    else:
        # 回退方案：使用 PowerShell 强制重启（某些系统/远程会话下不允许直接 stop）
        ps_cmd = [
            "powershell",
            "-NoProfile",
            "-ExecutionPolicy",
            "Bypass",
            "-Command",
            "Restart-Service -Name TermService -Force",
        ]
        try:
            start_proc = subprocess.run(ps_cmd, capture_output=True, text=True)
        except Exception:
            start_proc = None
    return stop_proc, start_proc


class RdpPortGUI:
    def __init__(self, root: tk.Tk):
        self.root = root
        self.root.title("更改远程桌面 (RDP) 端口")
        # 设定基础尺寸并支持水平拉伸，再居中显示
        self.base_width = 600
        self.base_height = 300
        self.root.geometry(f"{self.base_width}x{self.base_height}")
        self.root.minsize(self.base_width, self.base_height)
        self.root.resizable(True, False)

        # 统一字体以提升观感（使用 TkDefaultFont 配置更稳妥）
        try:
            df = tkfont.nametofont("TkDefaultFont")
            df.configure(family="Segoe UI", size=10)
            self.root.option_add("*Font", df)
        except Exception:
            pass

        self.current_port_var = tk.StringVar()
        self.new_port_var = tk.StringVar()
        # 更改后将自动重启远程桌面服务，无需开关

        self._build_ui()
        self._load_current_port()
        self._center_window(self.base_width, self.base_height)

    def _center_window(self, width: int, height: int):
        # 居中显示窗口
        self.root.update_idletasks()
        screen_w = self.root.winfo_screenwidth()
        screen_h = self.root.winfo_screenheight()
        x = (screen_w // 2) - (width // 2)
        y = (screen_h // 2) - (height // 2)
        self.root.geometry(f"{width}x{height}+{x}+{y}")

    def _build_ui(self):
        pad_opts = {"padx": 12, "pady": 8}

        frm = ttk.Frame(self.root)
        frm.pack(fill=tk.BOTH, expand=True)
        # 让两列按比例扩展，避免右侧控件被裁切
        frm.columnconfigure(0, weight=1)
        frm.columnconfigure(1, weight=1)
        # 增加一个可伸缩的空行，将底部版权推到更靠下
        frm.rowconfigure(2, weight=1)

        # 顶部主题选择
        # 主内容卡片
        card = ttk.Labelframe(frm, text="端口设置")
        card.grid(row=0, column=0, columnspan=2, sticky=tk.NSEW, padx=12, pady=(6, 12))
        card.columnconfigure(0, weight=1)
        card.columnconfigure(1, weight=1)

        ttk.Label(card, text="源端口（当前正在使用的端口，1-65535）：").grid(row=0, column=0, sticky=tk.W, **pad_opts)
        ttk.Entry(card, textvariable=self.current_port_var, width=12).grid(row=0, column=1, sticky=tk.EW, **pad_opts)

        ttk.Label(card, text="目标端口（我要改成的端口，1-65535）：").grid(row=1, column=0, sticky=tk.W, **pad_opts)
        ttk.Entry(card, textvariable=self.new_port_var, width=12).grid(row=1, column=1, sticky=tk.EW, **pad_opts)

        # 提示标签
        ttk.Label(card, text="更改后将自动重启远程桌面服务 (会断开会话)") \
            .grid(row=2, column=0, columnspan=2, sticky=tk.W, **pad_opts)

        # 操作按钮区
        btn_frame = ttk.Frame(card)
        btn_frame.grid(row=3, column=0, columnspan=2, sticky=tk.EW, padx=0, pady=(4, 8))
        ttk.Button(btn_frame, text="应用更改", command=self.on_apply).pack(side=tk.LEFT, padx=12, pady=6)
        ttk.Button(btn_frame, text="退出", command=self.root.destroy).pack(side=tk.RIGHT, padx=12, pady=6)

        note = (
            "提示：若无法自动读取当前端口，请手动输入当前端口；更改后使用 远程地址:端口 连接，例如 192.168.1.10:33890。"
        )
        ttk.Label(frm, text=note, foreground="#666", wraplength=self.base_width - 24) \
            .grid(row=1, column=0, columnspan=2, sticky=tk.W, padx=12)

        # 底部居中版权信息
        copyright_text = "版权所有：速光网络软件开发 suguang.cc 抖音号：dubaishun12"
        ttk.Label(frm, text=copyright_text, foreground="#888", anchor="center") \
            .grid(row=3, column=0, columnspan=2, sticky=tk.EW, padx=12, pady=(6, 12))

    # 主题固定为 xpnative，无需动态切换

    def _load_current_port(self):
        port = read_rdp_port()
        if port is None:
            # 未能确定端口时，允许用户手动输入，不设默认值
            self.current_port_var.set("")
            self.new_port_var.set("")
        else:
            self.current_port_var.set(str(port))
            self.new_port_var.set(str(port))

    def _validate_port(self, s: str) -> Tuple[bool, Optional[int]]:
        try:
            p = int(s)
        except ValueError:
            return False, None
        if 1 <= p <= 65535:
            return True, p
        return False, None

    def on_apply(self):
        ok_target, port = self._validate_port(self.new_port_var.get().strip())
        if not ok_target:
            messagebox.showerror("无效目标端口", "目标端口需为 1-65535 之间的整数")
            return

        ok_current, current_port_val = self._validate_port(self.current_port_var.get().strip())
        if self.current_port_var.get().strip() and not ok_current:
            messagebox.showerror("无效当前端口", "当前端口（可选）如填写需为 1-65535 之间的整数")
            return

        do_restart = True

        # 若非管理员，尝试提权并将操作作为参数传递
        if not is_admin():
            if messagebox.askyesno(
                "需要管理员权限",
                "写入注册表与管理防火墙/服务需要管理员权限。\n是否以管理员重新运行并应用更改？",
            ):
                args = [
                    "--apply-port", str(port),
                ]
                # 自动重启服务
                args.append("--restart-service")
                relaunch_as_admin(args)
                # 退出当前非管理员实例
                self.root.after(300, self.root.destroy)
            return

        # 管理员路径，直接执行
        try:
            write_rdp_port(port)
        except PermissionError:
            messagebox.showerror("权限不足", "写入注册表失败：请以管理员身份运行")
            return
        except Exception as e:
            messagebox.showerror("写入失败", f"写入端口失败：{e}")
            return

        def _explain(code: int) -> str:
            mapping = {
                0: "成功",
                1051: "可能存在依赖/会话阻止停止",
                1056: "服务已在运行",
                1058: "服务被禁用",
                1060: "服务不存在",
                1061: "当前无法接收控制",
                1062: "服务未启动",
                1053: "服务未及时响应",
            }
            return mapping.get(code, "未知")

        svc_msg = ""
        if do_restart:
            stop_proc, start_proc = restart_rdp_service()
            msg_parts = []
            msg_parts.append(f"stop: {stop_proc.returncode}({_explain(stop_proc.returncode)})")
            if start_proc is not None:
                msg_parts.append(f"start: {start_proc.returncode}({_explain(start_proc.returncode)})")
            else:
                msg_parts.append("start: N/A")
            svc_msg = "\n远程桌面服务重启结果：" + ", ".join(msg_parts)

        self.current_port_var.set(str(port))
        msg_head = (
            f"已将 RDP 端口{f'从 {current_port_val} ' if current_port_val else ''}更新为 {port}。"
        )
        messagebox.showinfo("已应用", msg_head + svc_msg)


def run_cli_actions(args: List[str]) -> bool:
    """在被提权重启时，处理命令行参数执行相应操作。返回是否执行过。"""
    if not args:
        return False
    if "--apply-port" in args:
        try:
            idx = args.index("--apply-port")
            port = int(args[idx + 1])
        except Exception:
            messagebox.showerror("参数错误", "--apply-port 后需跟合法端口号")
            return True

        try:
            write_rdp_port(port)
        except Exception as e:
            messagebox.showerror("写入失败", f"写入端口失败：{e}")
            return True

        msg = f"已将远程桌面端口更新为 {port}。"
        if "--restart-service" in args:
            stop_proc, start_proc = restart_rdp_service()
            def _explain(code: int) -> str:
                mapping = {
                    0: "成功",
                    1051: "可能存在依赖/会话阻止停止",
                    1056: "服务已在运行",
                    1058: "服务被禁用",
                    1060: "服务不存在",
                    1061: "当前无法接收控制",
                    1062: "服务未启动",
                    1053: "服务未及时响应",
                }
                return mapping.get(code, "未知")
            msg += "\n已尝试重启远程桌面服务。"
            msg += f" (stop={stop_proc.returncode}({_explain(stop_proc.returncode)}), start={getattr(start_proc, 'returncode', 'N/A')}{f'({_explain(start_proc.returncode)})' if start_proc else ''})"

        messagebox.showinfo("已应用", msg)
        return True

    return False


def main():
    # 如果通过参数要求直接操作（用于提权后的自动执行），先尝试执行
    if run_cli_actions(sys.argv[1:]):
        return

    # 高 DPI 适配（在支持的系统上）
    try:
        ctypes.windll.shcore.SetProcessDpiAwareness(1)
    except Exception:
        pass

    root = tk.Tk()
    # 固定使用 xpnative 主题（如不可用则忽略）
    try:
        s = ttk.Style()
        s.theme_use("xpnative")
    except Exception:
        pass
    RdpPortGUI(root)
    root.mainloop()


if __name__ == "__main__":
    main()